#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
:File       :train.py
:Description:
    载入数据与模型 -> 训练模型 -> 评估模型 

    大部分程序逻辑基于 参考链接中的实现，只不过 其使用的MNIST数据集并不是已经可视化之后的图片,而是原本的二进制数据
    所以,额外实现  已可视化的MNIST数据集 的载入逻辑
    Reference : https://gitee.com/kongfanhe/pytorch-tutorial/blob/master/test.py
:EditTime   :2025/05/03 19:33:28
:Author     :Kiumb
'''
import os 
import torch
from tqdm import tqdm
from model.fc_net import Net
from argparse import Namespace
import matplotlib.pyplot as plt
from dataloader.mnist import MNIST
from torch.utils.data import DataLoader


def parse_args():
    '''
    必要的参数配置函数
    '''
    cfg = Namespace(
        root = 'data/mnist',   # 数据集根目录
        exp  = 'exp',          # 实验目录 (保存每个轮次下的模型权重)

        device = 'cpu',        # 指定模型训练的设备  cpu  or  cuda 
        epoch  = 1,            # 训练轮次
        batch  = 16,           # 批量数
        
        lr     = 1e-3,         # 学习率
        weight_decay = 1e-4,   # 衰减项

    )
    return cfg

def evaluate(test_data, net,cfg):
    '''评估函数'''
    
    n_correct = 0
    n_total = 0
    with torch.no_grad():
        for (x, y) in test_data:
            x , y = x.to(cfg.device) , y.to(cfg.device)
            outputs = net(x.view(-1,28*28))  # Attention : 网络为MLP,需对输入数据进行展平操作
            for i, output in enumerate(outputs):
                if torch.argmax(output) == y[i]:
                    n_correct += 1
                n_total += 1
    return n_correct / n_total


def main():

    cfg = parse_args()
    checkpoints_path = os.path.join(cfg.exp,'checkpoints') # 保存模型权重
    eval_path = os.path.join(cfg.exp,'eval') # 保存模型在测试集上的预测结果
    os.makedirs(checkpoints_path,exist_ok=True)
    os.makedirs(eval_path,exist_ok=True)

    #---------------------------------#
    # 加载训练集和测试集
    #---------------------------------#
    train_dataset = MNIST(cfg.root,'train')
    test_dataset  = MNIST(cfg.root,'test')
    train_loader  = DataLoader(train_dataset,batch_size=cfg.batch,shuffle=True,pin_memory=True,drop_last=True)
    test_loader   = DataLoader(test_dataset,batch_size=cfg.batch,shuffle=False,pin_memory=True,drop_last=True)
    
    #---------------------------------#
    # 加载模型  | 未实现断点续训功能
    #---------------------------------#
    net = Net()
    net.to(cfg.device)

    #---------------------------------#
    # 加载 学习率调度器
    #---------------------------------#
    optimizer = torch.optim.Adam(net.parameters(), lr=cfg.lr,weight_decay=cfg.weight_decay)


    #---------------------------------#
    # 模型训练
    #---------------------------------#
    for epoch in range(cfg.epoch):
        for x,y in tqdm(train_loader,desc=f"Epoch:[{epoch+1}/{cfg.epoch}] in MNIST",unit='batch'):
            x , y = x.to(cfg.device) , y.to(cfg.device)
            net.zero_grad()
            output = net(x.view(-1,28*28))  # Attention : 网络为MLP,需对输入数据进行展平操作
            loss = torch.nn.functional.nll_loss(output, y)
            loss.backward()
            optimizer.step()
        print('Start Evaluation -> ',end='')
        acc = evaluate(test_loader,net,cfg)     # 未实现早停 early-stoping 策略
        print(f'ACC: {acc:.4f}')
        # 保存当前epoch下的模型权重
        torch.save(net.state_dict(),os.path.join(checkpoints_path,f'ep{epoch+1}-acc{acc:.2f}.pth'))

    #---------------------------------#
    # 模型评估 | 使用最后一个epoch的权重
    #---------------------------------#
    net.eval() # 切换至评估模式
    # 对测试集的一部分进行 预测 (预测结果同时保存在图像文件名中)
    for (n, (x, _)) in tqdm(enumerate(test_loader),desc=f"Evaluating MNIST Test Batches",unit='batch'):
        if n > 3:
            break
        x = x.to(cfg.device)
        predict = torch.argmax(net(x[0].view(-1,28*28)))  # Attention : 网络为MLP,需对输入数据进行展平操作
        predict = predict.to('cpu')
        # plt.figure(n)
        plt.imshow(x[0].view(28, 28).to('cpu'))
        plt.title("prediction: " + str(int(predict)))
        plt.savefig(eval_path+os.sep+f'p{predict}-{n}.jpg')
    # plt.show()

if __name__ == '__main__':
    main()
